import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from PIL import Image
from torch.backends import cudnn
from torch.utils.data import DataLoader
from tqdm import tqdm

from config import parser
from datalist import CartoonDataset
from model import UnetGenerator
from utils import slic, adaptive_slic, sscolor, \
    denormalize

SuperPixelDict = {
    'slic': slic,
    'adaptive_slic': adaptive_slic,
    'sscolor': sscolor}
from torchvision import transforms

data_transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
# resize2scale = transforms.Resize
infer_transform = transforms.Compose([
    transforms.Resize((512, 512)),
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
'''
这里的问题可能是出的引导图可能有问题
'''



'''
在运行train.py之前，首先需要训练pretrain.py
因为你会在运行train.py的时候发现，卧槽
怎么训练的generator这么差，时间上来说这可能也是不允许的，
所以为了保证generator的稳定，你可以先训练generator，保证它能实实在在的生成真实的图片
最后通过这个模型，去train.py里面把它逐渐变的卡通化
这就是一个trick
'''


class train(object):
    def __init__(self):
        self.args = parser.parse_args()
        print(f"-----------{self.args.project_name}-------------")

        use_cuda = self.args.use_cuda and torch.cuda.is_available()
        if use_cuda:
            torch.cuda.manual_seed(self.args.seed)
            torch.cuda.manual_seed_all(self.args.seed)
        else:
            torch.manual_seed(self.args.seed)

        self.device = torch.device('cuda' if use_cuda else 'cpu')
        # TODO
        # 转到UBUNTU后，这里的num_workers可以更改，扩大数据dataloader的线程数，使得gpu保持一个100%的工作状态
        # 加快模型的训练速度
        # windows下不支持这种设置
        train_kwargs = {'num_workers': 0, 'pin_memory': True} if use_cuda else {}

        '''
        构造DataLoader
        '''
        self.train_dataset = CartoonDataset(self.args.data_path, self.args.photo_path)
        self.train_dataloader = DataLoader(self.train_dataset, batch_size=self.args.train_batch_size, shuffle=False,
                                           **train_kwargs)

        '''
        构造Model
        '''
        self.generator = UnetGenerator().to(self.device)

        '''
        加载模型
        '''
        if self.args.resume == False:
            try:
                print("load the weight from pretrained-weight file")
                pretrained_dict = torch.load("weights/generator/best.pth", map_location=self.device)['model_state_dict']
                new_state_dict = {}
                for k, v in pretrained_dict.items():
                    new_state_dict[k[7:]] = v
                self.generator.load_state_dict(new_state_dict, True)

                del new_state_dict

                print("Finished to load the weight")
            except:
                print("can not load weight \n train the model from scratch")

        if use_cuda:
            self.generator = torch.nn.DataParallel(self.generator, device_ids=range(torch.cuda.device_count()))

            cudnn.benchmark = True

        '''
        构造loss目标函数
        选择优化器
        学习率变化选择
        '''

        self.l1_loss = nn.L1Loss()
        from torchvision.models import vgg16
        self.vgg=vgg16(pretrained=True).to(self.device).features[:22]


        self.optimizer_g = optim.Adam(self.generator.parameters(), lr=0.0001)

        self.scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(self.optimizer_g, T_max=5, eta_min=1e-5)

        for epoch in range(1, self.args.epochs + 1):
            self.train(epoch)
            if epoch % 1 == 0:
                self.test(epoch)

        torch.cuda.empty_cache()
        print("finish model training")

    def train(self, epoch):
        self.generator.train()
        g_loss=[]

        pbar = tqdm(self.train_dataloader, desc=f'Train Epoch{epoch}/{self.args.epochs}')
        for cartoon, photo in pbar:
            #################################################################################

            '''
            生成网络训练
            '''

            self.optimizer_g.zero_grad()
            img = photo.to(self.device)

            generator_img = self.generator(img)  # generator


            self.get_image(generator_img)
            _,c,h,w=generator_img.shape
            generator_img=self.vgg(generator_img)/(c*h*w)
            img=self.vgg(img)/(c*h*w)
            loss = self.l1_loss(generator_img, img)
            loss.backward()
            g_loss.append(loss.item())
            self.optimizer_g.step()

            pbar.set_description(
                f'Train Epoch:{epoch}/{self.args.epochs} g_loss :{np.average(g_loss)} ')
        self.scheduler.step()

    def test(self, epoch):
        # self.generator.eval()

        with torch.no_grad():
            img = Image.open(r"G:\datasets\VOC\VOC2012\JPEGImages\2007_000175.jpg")
            img = data_transform(img).unsqueeze(0)
            output = self.generator(img)

            # output=output.cpu().numpy()
            #
            # output=output[0]
            output = transforms.ToPILImage()(output[0])
            # output=Image.fromarray(output).convert('RGB')
            output.save("./results/" + str(epoch) + '.jpg')

            torch.save({
                'model_state_dict': self.generator.state_dict()
            },
                './weights/generator/best.pth')

            print("save models")

    def get_image(self, image):
        # image=image[0]
        output = (denormalize(image.permute((0, 2, 3, 1)).
                              detach().
                              to('cpu').
                              numpy()) * 255).astype('uint8')
        output = output[0]

        output = Image.fromarray(output).convert('RGB')
        output.save("./results/1.jpg")


train = train()
